from torch import nn
import torch


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(3, 16, 3),
            nn.ReLU(),
            nn.MaxPool2d(3),
            nn.Conv2d(16, 22, 3),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(22, 32, 3),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 42, 3),
            nn.ReLU(),
            nn.Conv2d(42, 64, 3),
            nn.ReLU()
        )
        self.out_layer = nn.Sequential(
            nn.Linear(64 * 19 * 19, 26),
        )

    def forward(self, x):
        out = self.layers(x).reshape(-1, 64 * 19 * 19)
        out = self.out_layer(out)
        label = out[:, 0]
        position = out[:, 1:5]
        sort = out[:, 5:]
        return label, position, sort


class Net_v2(nn.Module):
    def __init__(self):
        super(Net_v2, self).__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(3, 16, 3),
            nn.ReLU(),
            nn.MaxPool2d(3),
            nn.Conv2d(16, 22, 3),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(22, 32, 5),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 82, 3),
            nn.ReLU(),
            nn.Conv2d(82, 128, 3),
            nn.ReLU(),
            nn.Conv2d(128, 25, 3),
            nn.ReLU()
        )
        self.label_layer = nn.Sequential(
            nn.Conv2d(25, 1, 3),
            nn.ReLU(),

        )
        self.position_layers = nn.Sequential(
            nn.Conv2d(25, 4, 3),
            nn.ReLU()
        )
        self.sort_layers = nn.Sequential(
            nn.Conv2d(25, 20, 3),
            nn.ReLU()
        )

    def forward(self, x):
        out = self.layers(x)
        #print(out)
        label = self.label_layer(out)
        label = torch.squeeze(label, dim=2)
        label = torch.squeeze(label, dim=2)
        label = torch.squeeze(label, dim=1)
        position = self.position_layers(out)
        position = torch.squeeze(position, dim=2)
        position = torch.squeeze(position, dim=2)
        sort = self.sort_layers(out)
        sort = torch.squeeze(sort, dim=2)
        sort = torch.squeeze(sort, dim=2)
        return label, position, sort


# class Net_v3(nn.Module):
#     def __init__(self):
#         super(Net_v3, self).__init__()
#         self.layers = nn.Sequential(
#             nn.Conv2d(3, 16, 3),
#             nn.ReLU(),
#             nn.MaxPool2d(3),
#             nn.Conv2d(16, 32, 5),
#             nn.ReLU(),
#             nn.MaxPool2d(2),
#             nn.Conv2d(32, 64, 5),
#             nn.ReLU(),
#             nn.MaxPool2d(2),
#             nn.Conv2d(64, 128, 3),
#             nn.ReLU(),
#             nn.MaxPool2d(2),
#             nn.Conv2d(128, 132, 3),
#             nn.ReLU(),
#             nn.Conv2d(132,25,7),
#             nn.ReLU()
#         )
#         self.position_layers=
#     def forward(self, x):
#         out = torch.squeeze(self.layers(x), dim=2)
#         out = torch.squeeze(out, dim=2)
#
#         return position, sort


if __name__ == '__main__':
    net = Net_v2()
    x = torch.randn(1, 3, 300, 300)
    proce = net(x)
    print(net(x)[0].shape)
